/*
 * Copyright 2012-2020 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.example.gateway.balancer;
import java.net.URI;
import java.util.List;
import com.example.common.entity.BaseUserInfo;
import com.example.common.util.RequestUtil;
import com.example.gateway.cluster.ClusterManager;
import com.example.gateway.cluster.ServerNode;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.cloud.client.loadbalancer.*;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.SelectedInstanceCallback;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;

/**
 * A random-based implementation of {@link ReactorServiceInstanceLoadBalancer}.
 * 默认的 负载均衡过滤器
 * ReactiveLoadBalancerClientFilter 类的 choose方法 会调用此类的choose方法
 * private Mono<Response<ServiceInstance>> choose(Request<RequestDataContext> lbRequest, String serviceId,
 * Set<LoadBalancerLifecycle> supportedLifecycleProcessors) {
 * //可以从clientFactory中拿到我们自定义的负载均衡策略实例
 * ReactorLoadBalancer<ServiceInstance> loadBalancer = this.clientFactory.getInstance(serviceId,
 * ReactorServiceInstanceLoadBalancer.class);
 * if (loadBalancer == null) {
 * throw new NotFoundException("No loadbalancer available for " + serviceId);
 * }
 * supportedLifecycleProcessors.forEach(lifecycle -> lifecycle.onStart(lbRequest));
 * //重点
 * return loadBalancer.choose(lbRequest);*
 * }
 *
 * @author 谢飞机
 * @date 2021-08-25
 * 自定义负载均衡策略
 */
public class CustomLoadBalancerStrategy implements ReactorServiceInstanceLoadBalancer {

    private static final Log log = LogFactory.getLog(CustomLoadBalancerStrategy.class);

    private final String serviceId;

    private ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;

    private ClusterManager clusterManager;

    /**
     * @param serviceInstanceListSupplierProvider a provider of
     *                                            {@link ServiceInstanceListSupplier} that will be used to get available instances
     * @param serviceId                           id of the service for which to choose an instance
     */
    public CustomLoadBalancerStrategy(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider,
                                      String serviceId,
                                      ClusterManager clusterManager) {
        this.serviceId = serviceId;
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
        this.clusterManager = clusterManager;
    }

    @SuppressWarnings("rawtypes")
    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {

        ServiceInstanceListSupplier supplier = serviceInstanceListSupplierProvider
                .getIfAvailable(NoopServiceInstanceListSupplier::new);

        Flux<List<ServiceInstance>> listFlux = supplier.get(request);

        return listFlux.next()
                .map(serviceInstances -> processInstanceResponse(request, supplier, serviceInstances));
    }

    private Response<ServiceInstance> processInstanceResponse(Request request, ServiceInstanceListSupplier supplier,
                                                              List<ServiceInstance> serviceInstances) {
        Response<ServiceInstance> serviceInstanceResponse = getInstanceResponse(request, serviceInstances);

        boolean isHasServer = serviceInstanceResponse.hasServer();
        if (supplier instanceof SelectedInstanceCallback && isHasServer) {
            ((SelectedInstanceCallback) supplier).selectedServiceInstance(serviceInstanceResponse.getServer());
        }
        return serviceInstanceResponse;
    }

    /**
     * 自定义负载均衡规则：通过用户Id 在哈希环上找到ws 实例的ip
     * 客户端和服务端建立连接时才会根据负载均衡策略进行实例的选择
     *
     * @param instances
     * @return
     */
    private Response<ServiceInstance> getInstanceResponse(Request request, List<ServiceInstance> instances) {
        if (instances.isEmpty()) {
            if (log.isWarnEnabled()) {
                log.warn("- - - - - - - - - No servers available for service - - - - - - - - -  : " + serviceId);
            }
            return new EmptyResponse();
        }

        DefaultRequestContext defaultRequestContext = (DefaultRequestContext) request.getContext();

        RequestData requestData = (RequestData) defaultRequestContext.getClientRequest();

        URI uri = requestData.getUrl();

        BaseUserInfo baseUserInfo = RequestUtil.getBaseUserInfor(uri);

        /**
         * 普通轮询负载均衡
         * int index = ThreadLocalRandom.current().nextInt(instances.size());
         * 自定义负载均衡
         * 用户Id的hashcode % 服务实例数
         */

        int index = baseUserInfo.getId().hashCode() % instances.size();

        ServiceInstance instance = instances.get(index);

        //添加节点和服务实例之间的映射关系
        ServerNode serverNode = new ServerNode();
        serverNode.addClientNode(uri.getHost(), uri.getPort(), baseUserInfo.getId(),instance.getHost());
        clusterManager.getHashLoopMap().put(instance.getHost(),serverNode);
        return new DefaultResponse(instance);
    }

}
